#%% 定义神经网络
import torch.nn as nn
import torch.nn.functional as F 

# 一个更简单的模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        
        # 卷积1操作，1个通道，10个卷积核，卷积核大小5*5
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        
        # 卷积2操作
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        
        # 最大池化
        self.pooling = nn.MaxPool2d(2)
        
        # 全连接
        self.fc = nn.Linear(320, 10)
 
 
    def forward(self, x):
        # flatten data from (n,1,28,28) to (n, 784)
        batch_size = x.size(0)
        x = F.relu(self.pooling(self.conv1(x)))
        x = F.relu(self.pooling(self.conv2(x)))
        x = x.view(batch_size, -1) # -1 此处自动算出的是320
        x = self.fc(x)
 
        return x